def __init__(self, dims_in, init_data=None):
    super().__init__()
    self.dims_in = dims_in[0]
    param_dims = [1, self.dims_in[0]] + [1 for i in range(len(self.dims_in) - 1)]
    self.scale = nn.Parameter(torch.zeros(*param_dims))
    self.bias = nn.Parameter(torch.zeros(*param_dims))
    if init_data:
        self.initialize_with_data(init_data)
    else:
        self.init_on_next_batch = True

def add_module(self, name, module):
    r"""Adds a child module to the current module.
    The module can be accessed as an attribute using the given name.
    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\"")
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module

def apply(self, fn):
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`torch-nn-init`).
    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule
    Returns:
        Module: self
    Example::
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.data.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self

def buffers(self, recurse=True):
    r"""Returns an iterator over module buffers.
    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.
    Yields:
        torch.Tensor: module buffer
    Example::
        >>> for buf in model.buffers():
        >>>     print(type(buf.data), buf.size())
        <class 'torch.FloatTensor'> (20L,)
        <class 'torch.FloatTensor'> (20L, 1L, 5L, 5L)
    """
    for name, buf in self.named_buffers(recurse=recurse):
        yield buf

def children(self):
    r"""Returns an iterator over immediate children modules.
    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module

def cpu(self):
    r"""Moves all model parameters and buffers to the CPU.
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())

def cuda(self, device=None):
    r"""Moves all model parameters and buffers to the GPU.
    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.
    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))

def double(self):
    r"""Casts all floating point parameters and buffers to ``double`` datatype.
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)

def eval(self):
    r"""Sets the module in evaluation mode.
    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.
    """
    return self.train(False)

def extra_repr(self):
    r"""Set the extra representation of the module
    To print customized extra information, you should reimplement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''

def float(self):
    r"""Casts all floating point parameters and buffers to float datatype.
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)

def forward(self, x, rev=False):
    if self.init_on_next_batch:
        self.initialize_with_data(x[0])
    if not rev:
        return [x[0] * self.scale.exp() + self.bias]
    else:
        return [(x[0] - self.bias) / self.scale.exp()]

def half(self):
    r"""Casts all floating point parameters and buffers to ``half`` datatype.
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)

def initialize_with_data(self, data):
    # Initialize to mean 0 and std 1 with sample batch
    # 'data' expected to be of shape (batch, channels[, ...])
    assert all([data.shape[i+1] == self.dims_in[i] for i in range(len(self.dims_in))]),\
        "Can't initialize ActNorm layer, provided data don't match input dimensions."
    self.scale.data.view(-1)[:] \
        = torch.log(1 / data.transpose(0,1).contiguous().view(self.dims_in[0], -1).std(dim=-1))
    data = data * self.scale.exp()
    self.bias.data.view(-1)[:] \
        = -data.transpose(0,1).contiguous().view(self.dims_in[0], -1).mean(dim=-1)
    self.init_on_next_batch = False

def jacobian(self, x, rev=False):
    if not rev:
        return (self.scale.sum() * np.prod(self.dims_in[1:])).repeat(x[0].shape[0])
    else:
        return -self.jacobian(x)

def load_state_dict(self, state_dict, strict=True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.
    Arguments:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys
    """
    missing_keys = []
    unexpected_keys = []
    error_msgs = []
    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        state_dict._metadata = metadata
    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')
    load(self)
    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))
    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)

def modules(self):
    r"""Returns an iterator over all modules in the network.
    Yields:
        Module: a module in the network
    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.
    Example::
        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)
        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)
    """
    for name, module in self.named_modules():
        yield module

def named_buffers(self, prefix='', recurse=True):
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.
    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.
    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer
    Example::
        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())
    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem

def named_children(self):
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.
    Yields:
        (string, Module): Tuple containing a name and child module
    Example::
        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)
    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module

def named_modules(self, memo=None, prefix=''):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.
    Yields:
        (string, Module): Tuple of name and module
    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.
    Example::
        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)
        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
    """
    if memo is None:
        memo = set()
    if self not in memo:
        memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix):
                yield m

def named_parameters(self, prefix='', recurse=True):
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.
    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.
    Yields:
        (string, Parameter): Tuple containing the name and parameter
    Example::
        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())
    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem

def output_dims(self, input_dims):
    assert len(input_dims) == 1, "Can only use 1 input"
    return input_dims

def parameters(self, recurse=True):
    r"""Returns an iterator over module parameters.
    This is typically passed to an optimizer.
    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.
    Yields:
        Parameter: module parameter
    Example::
        >>> for param in model.parameters():
        >>>     print(type(param.data), param.size())
        <class 'torch.FloatTensor'> (20L,)
        <class 'torch.FloatTensor'> (20L, 1L, 5L, 5L)
    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param

def register_backward_hook(self, hook):
    r"""Registers a backward hook on the module.
    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::
        hook(module, grad_input, grad_output) -> Tensor or None
    The :attr:`grad_input` and :attr:`grad_output` may be tuples if the
    module has multiple inputs or outputs. The hook should not modify its
    arguments, but it can optionally return a new gradient with respect to
    input that will be used in place of :attr:`grad_input` in subsequent
    computations.
    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    .. warning ::
        The current implementation will not have the presented behavior
        for complex :class:`Module` that perform many operations.
        In some failure cases, :attr:`grad_input` and :attr:`grad_output` will only
        contain the gradients for a subset of the inputs and outputs.
        For such :class:`Module`, you should use :func:`torch.Tensor.register_hook`
        directly on a specific input or output to get the required gradients.
    """
    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle

def register_buffer(self, name, tensor):
    r"""Adds a persistent buffer to the module.
    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the persistent state.
    Buffers can be accessed as attributes using given names.
    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor): buffer to be registered.
    Example::
        >>> self.register_buffer('running_mean', torch.zeros(num_features))
    """
    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor

def register_forward_hook(self, hook):
    r"""Registers a forward hook on the module.
    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::
        hook(module, input, output) -> None
    The hook should not modify the input or output.
    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle

def register_forward_pre_hook(self, hook):
    r"""Registers a forward pre-hook on the module.
    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::
        hook(module, input) -> None
    The hook should not modify the input.
    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle

def register_parameter(self, name, param):
    r"""Adds a parameter to the module.
    The parameter can be accessed as an attribute using given name.
    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter): parameter to be added to the module.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))
    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param

def share_memory(self):
    return self._apply(lambda t: t.share_memory_())

def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.
    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Returns:
        dict:
            a dictionary containing a whole state of the module
    Example::
        >>> module.state_dict().keys()
        ['bias', 'weight']
    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    for name, param in self._parameters.items():
        if param is not None:
            destination[prefix + name] = param if keep_vars else param.data
    for name, buf in self._buffers.items():
        if buf is not None:
            destination[prefix + name] = buf if keep_vars else buf.data
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination

def to(self, *args, **kwargs):
    r"""Moves and/or casts the parameters and buffers.
    This can be called as
    .. function:: to(device=None, dtype=None, non_blocking=False)
    .. function:: to(dtype, non_blocking=False)
    .. function:: to(tensor, non_blocking=False)
    Its signature is similar to :meth:`torch.Tensor.to`, but only accepts
    floating point desired :attr:`dtype` s. In addition, this method will
    only cast the floating point parameters and buffers to :attr:`dtype`
    (if given). The integral parameters and buffers will be moved
    :attr:`device`, if that is given, but with dtypes unchanged. When
    :attr:`non_blocking` is set, it tries to convert/move asynchronously
    with respect to the host if possible, e.g., moving CPU Tensors with
    pinned memory to CUDA devices.
    See below for examples.
    .. note::
        This method modifies the module in-place.
    Args:
        device (:class:`torch.device`): the desired device of the parameters
            and buffers in this module
        dtype (:class:`torch.dtype`): the desired floating point type of
            the floating point parameters and buffers in this module
        tensor (torch.Tensor): Tensor whose dtype and device are the desired
            dtype and device for all parameters and buffers in this module
    Returns:
        Module: self
    Example::
        >>> linear = nn.Linear(2, 2)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]])
        >>> linear.to(torch.double)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1913, -0.3420],
                [-0.5113, -0.2325]], dtype=torch.float64)
        >>> gpu1 = torch.device("cuda:1")
        >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
        >>> cpu = torch.device("cpu")
        >>> linear.to(cpu)
        Linear(in_features=2, out_features=2, bias=True)
        >>> linear.weight
        Parameter containing:
        tensor([[ 0.1914, -0.3420],
                [-0.5112, -0.2324]], dtype=torch.float16)
    """
    device, dtype, non_blocking = torch._C._nn._parse_to(*args, **kwargs)
    if dtype is not None:
        if not dtype.is_floating_point:
            raise TypeError('nn.Module.to only accepts floating point '
                            'dtypes, but got desired dtype={}'.format(dtype))
    def convert(t):
        return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
    return self._apply(convert)

def train(self, mode=True):
    r"""Sets the module in training mode.
    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.
    Returns:
        Module: self
    """
    self.training = mode
    for module in self.children():
        module.train(mode)
    return self

def type(self, dst_type):
    r"""Casts all parameters and buffers to :attr:`dst_type`.
    Arguments:
        dst_type (type or string): the desired type
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.type(dst_type))

def zero_grad(self):
    r"""Sets gradients of all model parameters to zero."""
    for p in self.parameters():
        if p.grad is not None:
            p.grad.detach_()
            p.grad.zero_()

def __init__(self, dims_in, dims_c=[], F_class=F_fully_connected, F_args={}, **kwargs):
    warnings.warn(F"{self.__class__.__name__} is deprecated and will be removed in the public release. "
                  F"Use {orig_class.__name__} instead.",
                  DeprecationWarning)
    def coeff_func_wrapper(ch_in, ch_out):
        return F_class(ch_in, ch_out, **F_args)
    super().__init__(dims_in, dims_c, subnet_constructor=coeff_func_wrapper, **kwargs)

def add_module(self, name, module):
    r"""Adds a child module to the current module.
    The module can be accessed as an attribute using the given name.
    Args:
        name (string): name of the child module. The child module can be
            accessed from this module using the given name
        module (Module): child module to be added to the module.
    """
    if not isinstance(module, Module) and module is not None:
        raise TypeError("{} is not a Module subclass".format(
            torch.typename(module)))
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("module name should be a string. Got {}".format(
            torch.typename(name)))
    elif hasattr(self, name) and name not in self._modules:
        raise KeyError("attribute '{}' already exists".format(name))
    elif '.' in name:
        raise KeyError("module name can't contain \".\"")
    elif name == '':
        raise KeyError("module name can't be empty string \"\"")
    self._modules[name] = module

def apply(self, fn):
    r"""Applies ``fn`` recursively to every submodule (as returned by ``.children()``)
    as well as self. Typical use includes initializing the parameters of a model
    (see also :ref:`torch-nn-init`).
    Args:
        fn (:class:`Module` -> None): function to be applied to each submodule
    Returns:
        Module: self
    Example::
        >>> def init_weights(m):
        >>>     print(m)
        >>>     if type(m) == nn.Linear:
        >>>         m.weight.data.fill_(1.0)
        >>>         print(m.weight)
        >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
        >>> net.apply(init_weights)
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Linear(in_features=2, out_features=2, bias=True)
        Parameter containing:
        tensor([[ 1.,  1.],
                [ 1.,  1.]])
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
    """
    for module in self.children():
        module.apply(fn)
    fn(self)
    return self

def buffers(self, recurse=True):
    r"""Returns an iterator over module buffers.
    Args:
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.
    Yields:
        torch.Tensor: module buffer
    Example::
        >>> for buf in model.buffers():
        >>>     print(type(buf.data), buf.size())
        <class 'torch.FloatTensor'> (20L,)
        <class 'torch.FloatTensor'> (20L, 1L, 5L, 5L)
    """
    for name, buf in self.named_buffers(recurse=recurse):
        yield buf

def children(self):
    r"""Returns an iterator over immediate children modules.
    Yields:
        Module: a child module
    """
    for name, module in self.named_children():
        yield module

def cpu(self):
    r"""Moves all model parameters and buffers to the CPU.
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cpu())

def cuda(self, device=None):
    r"""Moves all model parameters and buffers to the GPU.
    This also makes associated parameters and buffers different objects. So
    it should be called before constructing optimizer if the module will
    live on GPU while being optimized.
    Arguments:
        device (int, optional): if specified, all parameters will be
            copied to that device
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.cuda(device))

def double(self):
    r"""Casts all floating point parameters and buffers to ``double`` datatype.
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.double() if t.is_floating_point() else t)

def e(self, s):
    return torch.exp(self.clamp * 0.636 * torch.atan(s))

def eval(self):
    r"""Sets the module in evaluation mode.
    This has any effect only on certain modules. See documentations of
    particular modules for details of their behaviors in training/evaluation
    mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
    etc.
    """
    return self.train(False)

def extra_repr(self):
    r"""Set the extra representation of the module
    To print customized extra information, you should reimplement
    this method in your own modules. Both single-line and multi-line
    strings are acceptable.
    """
    return ''

def float(self):
    r"""Casts all floating point parameters and buffers to float datatype.
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.float() if t.is_floating_point() else t)

def forward(self, x, c=[], rev=False):
    x1, x2 = torch.split(x[0], [self.split_idx, self.channels - self.split_idx], dim=1)
    if not rev:
        x1_c = torch.cat([x1, *c], 1) if self.conditional else x1
        s, t = self.s(x1_c), self.t(x1_c)
        y2 = self.e(s) * x2 + t
        self.last_s = s
        return [torch.cat((x1, y2), 1)]
    else:
        y1, y2 = x1, x2
        y1_c = torch.cat([y1, *c], 1) if self.conditional else y1
        s, t = self.s(y1_c), self.t(y1_c)
        x2 = (y2 - t) / self.e(s)
        self.last_s = s
        return [torch.cat((y1, x2), 1)]

def half(self):
    r"""Casts all floating point parameters and buffers to ``half`` datatype.
    Returns:
        Module: self
    """
    return self._apply(lambda t: t.half() if t.is_floating_point() else t)

def jacobian(self, x, c=[], rev=False):
    x1, x2 = torch.split(x[0], [self.split_idx, self.channels - self.split_idx], dim=1)
    if not rev:
        jac = self.log_e(self.last_s)
    else:
        jac = -self.log_e(self.last_s)
    return torch.sum(jac, dim=tuple(range(1, self.ndims+1)))

def load_state_dict(self, state_dict, strict=True):
    r"""Copies parameters and buffers from :attr:`state_dict` into
    this module and its descendants. If :attr:`strict` is ``True``, then
    the keys of :attr:`state_dict` must exactly match the keys returned
    by this module's :meth:`~torch.nn.Module.state_dict` function.
    Arguments:
        state_dict (dict): a dict containing parameters and
            persistent buffers.
        strict (bool, optional): whether to strictly enforce that the keys
            in :attr:`state_dict` match the keys returned by this module's
            :meth:`~torch.nn.Module.state_dict` function. Default: ``True``
    Returns:
        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
            * **missing_keys** is a list of str containing the missing keys
            * **unexpected_keys** is a list of str containing the unexpected keys
    """
    missing_keys = []
    unexpected_keys = []
    error_msgs = []
    # copy state_dict so _load_from_state_dict can modify it
    metadata = getattr(state_dict, '_metadata', None)
    state_dict = state_dict.copy()
    if metadata is not None:
        state_dict._metadata = metadata
    def load(module, prefix=''):
        local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
        module._load_from_state_dict(
            state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs)
        for name, child in module._modules.items():
            if child is not None:
                load(child, prefix + name + '.')
    load(self)
    if strict:
        if len(unexpected_keys) > 0:
            error_msgs.insert(
                0, 'Unexpected key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in unexpected_keys)))
        if len(missing_keys) > 0:
            error_msgs.insert(
                0, 'Missing key(s) in state_dict: {}. '.format(
                    ', '.join('"{}"'.format(k) for k in missing_keys)))
    if len(error_msgs) > 0:
        raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
                           self.__class__.__name__, "\n\t".join(error_msgs)))
    return _IncompatibleKeys(missing_keys, unexpected_keys)

def log_e(self, s):
    '''log of the nonlinear function e'''
    return self.clamp * 0.636 * torch.atan(s)

def modules(self):
    r"""Returns an iterator over all modules in the network.
    Yields:
        Module: a module in the network
    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.
    Example::
        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.modules()):
                print(idx, '->', m)
        0 -> Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        )
        1 -> Linear(in_features=2, out_features=2, bias=True)
    """
    for name, module in self.named_modules():
        yield module

def named_buffers(self, prefix='', recurse=True):
    r"""Returns an iterator over module buffers, yielding both the
    name of the buffer as well as the buffer itself.
    Args:
        prefix (str): prefix to prepend to all buffer names.
        recurse (bool): if True, then yields buffers of this module
            and all submodules. Otherwise, yields only buffers that
            are direct members of this module.
    Yields:
        (string, torch.Tensor): Tuple containing the name and buffer
    Example::
        >>> for name, buf in self.named_buffers():
        >>>    if name in ['running_var']:
        >>>        print(buf.size())
    """
    gen = self._named_members(
        lambda module: module._buffers.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem

def named_children(self):
    r"""Returns an iterator over immediate children modules, yielding both
    the name of the module as well as the module itself.
    Yields:
        (string, Module): Tuple containing a name and child module
    Example::
        >>> for name, module in model.named_children():
        >>>     if name in ['conv4', 'conv5']:
        >>>         print(module)
    """
    memo = set()
    for name, module in self._modules.items():
        if module is not None and module not in memo:
            memo.add(module)
            yield name, module

def named_modules(self, memo=None, prefix=''):
    r"""Returns an iterator over all modules in the network, yielding
    both the name of the module as well as the module itself.
    Yields:
        (string, Module): Tuple of name and module
    Note:
        Duplicate modules are returned only once. In the following
        example, ``l`` will be returned only once.
    Example::
        >>> l = nn.Linear(2, 2)
        >>> net = nn.Sequential(l, l)
        >>> for idx, m in enumerate(net.named_modules()):
                print(idx, '->', m)
        0 -> ('', Sequential(
          (0): Linear(in_features=2, out_features=2, bias=True)
          (1): Linear(in_features=2, out_features=2, bias=True)
        ))
        1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
    """
    if memo is None:
        memo = set()
    if self not in memo:
        memo.add(self)
        yield prefix, self
        for name, module in self._modules.items():
            if module is None:
                continue
            submodule_prefix = prefix + ('.' if prefix else '') + name
            for m in module.named_modules(memo, submodule_prefix):
                yield m

def named_parameters(self, prefix='', recurse=True):
    r"""Returns an iterator over module parameters, yielding both the
    name of the parameter as well as the parameter itself.
    Args:
        prefix (str): prefix to prepend to all parameter names.
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.
    Yields:
        (string, Parameter): Tuple containing the name and parameter
    Example::
        >>> for name, param in self.named_parameters():
        >>>    if name in ['bias']:
        >>>        print(param.size())
    """
    gen = self._named_members(
        lambda module: module._parameters.items(),
        prefix=prefix, recurse=recurse)
    for elem in gen:
        yield elem

def output_dims(self, input_dims):
    assert len(input_dims) == 1, "Can only use one input."
    return input_dims

def parameters(self, recurse=True):
    r"""Returns an iterator over module parameters.
    This is typically passed to an optimizer.
    Args:
        recurse (bool): if True, then yields parameters of this module
            and all submodules. Otherwise, yields only parameters that
            are direct members of this module.
    Yields:
        Parameter: module parameter
    Example::
        >>> for param in model.parameters():
        >>>     print(type(param.data), param.size())
        <class 'torch.FloatTensor'> (20L,)
        <class 'torch.FloatTensor'> (20L, 1L, 5L, 5L)
    """
    for name, param in self.named_parameters(recurse=recurse):
        yield param

def register_backward_hook(self, hook):
    r"""Registers a backward hook on the module.
    The hook will be called every time the gradients with respect to module
    inputs are computed. The hook should have the following signature::
        hook(module, grad_input, grad_output) -> Tensor or None
    The :attr:`grad_input` and :attr:`grad_output` may be tuples if the
    module has multiple inputs or outputs. The hook should not modify its
    arguments, but it can optionally return a new gradient with respect to
    input that will be used in place of :attr:`grad_input` in subsequent
    computations.
    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    .. warning ::
        The current implementation will not have the presented behavior
        for complex :class:`Module` that perform many operations.
        In some failure cases, :attr:`grad_input` and :attr:`grad_output` will only
        contain the gradients for a subset of the inputs and outputs.
        For such :class:`Module`, you should use :func:`torch.Tensor.register_hook`
        directly on a specific input or output to get the required gradients.
    """
    handle = hooks.RemovableHandle(self._backward_hooks)
    self._backward_hooks[handle.id] = hook
    return handle

def register_buffer(self, name, tensor):
    r"""Adds a persistent buffer to the module.
    This is typically used to register a buffer that should not to be
    considered a model parameter. For example, BatchNorm's ``running_mean``
    is not a parameter, but is part of the persistent state.
    Buffers can be accessed as attributes using given names.
    Args:
        name (string): name of the buffer. The buffer can be accessed
            from this module using the given name
        tensor (Tensor): buffer to be registered.
    Example::
        >>> self.register_buffer('running_mean', torch.zeros(num_features))
    """
    if '_buffers' not in self.__dict__:
        raise AttributeError(
            "cannot assign buffer before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("buffer name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("buffer name can't contain \".\"")
    elif name == '':
        raise KeyError("buffer name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._buffers:
        raise KeyError("attribute '{}' already exists".format(name))
    elif tensor is not None and not isinstance(tensor, torch.Tensor):
        raise TypeError("cannot assign '{}' object to buffer '{}' "
                        "(torch Tensor or None required)"
                        .format(torch.typename(tensor), name))
    else:
        self._buffers[name] = tensor

def register_forward_hook(self, hook):
    r"""Registers a forward hook on the module.
    The hook will be called every time after :func:`forward` has computed an output.
    It should have the following signature::
        hook(module, input, output) -> None
    The hook should not modify the input or output.
    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_hooks)
    self._forward_hooks[handle.id] = hook
    return handle

def register_forward_pre_hook(self, hook):
    r"""Registers a forward pre-hook on the module.
    The hook will be called every time before :func:`forward` is invoked.
    It should have the following signature::
        hook(module, input) -> None
    The hook should not modify the input.
    Returns:
        :class:`torch.utils.hooks.RemovableHandle`:
            a handle that can be used to remove the added hook by calling
            ``handle.remove()``
    """
    handle = hooks.RemovableHandle(self._forward_pre_hooks)
    self._forward_pre_hooks[handle.id] = hook
    return handle

def register_parameter(self, name, param):
    r"""Adds a parameter to the module.
    The parameter can be accessed as an attribute using given name.
    Args:
        name (string): name of the parameter. The parameter can be accessed
            from this module using the given name
        param (Parameter): parameter to be added to the module.
    """
    if '_parameters' not in self.__dict__:
        raise AttributeError(
            "cannot assign parameter before Module.__init__() call")
    elif not isinstance(name, torch._six.string_classes):
        raise TypeError("parameter name should be a string. "
                        "Got {}".format(torch.typename(name)))
    elif '.' in name:
        raise KeyError("parameter name can't contain \".\"")
    elif name == '':
        raise KeyError("parameter name can't be empty string \"\"")
    elif hasattr(self, name) and name not in self._parameters:
        raise KeyError("attribute '{}' already exists".format(name))
    if param is None:
        self._parameters[name] = None
    elif not isinstance(param, Parameter):
        raise TypeError("cannot assign '{}' object to parameter '{}' "
                        "(torch.nn.Parameter or None required)"
                        .format(torch.typename(param), name))
    elif param.grad_fn:
        raise ValueError(
            "Cannot assign non-leaf Tensor to parameter '{0}'. Model "
            "parameters must be created explicitly. To express '{0}' "
            "as a function of another Tensor, compute the value in "
            "the forward() method.".format(name))
    else:
        self._parameters[name] = param

def share_memory(self):
    return self._apply(lambda t: t.share_memory_())

def state_dict(self, destination=None, prefix='', keep_vars=False):
    r"""Returns a dictionary containing a whole state of the module.
    Both parameters and persistent buffers (e.g. running averages) are
    included. Keys are corresponding parameter and buffer names.
    Returns:
        dict:
            a dictionary containing a whole state of the module
    Example::
        >>> module.state_dict().keys()
        ['bias', 'weight']
    """
    if destination is None:
        destination = OrderedDict()
        destination._metadata = OrderedDict()
    destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
    for name, param in self._parameters.items():
        if param is not None:
            destination[prefix + name] = param if keep_vars else param.data
    for name, buf in self._buffers.items():
        if buf is not None:
            destination[prefix + name] = buf if keep_vars else buf.data
    for name, module in self._modules.items():
        if module is not None:
            module.state_dict(destination, prefix + name + '.', keep_vars=keep_vars)
    for hook in self._state_dict_hooks.values():
        hook_result = hook(self, destination, prefix, local_metadata)
        if hook_result is not None:
            destination = hook_result
    return destination